#include "dag.h"

#include <set>
#include <map>
#include <vector>
#include <algorithm>

using namespace dag;
using namespace std;

static bool equalNodes(const Edge* edge1, const Edge* edge2);

Graph::Graph()
:	m_HeadNode(NULL),
	m_TailNode(NULL)
{
}

Graph::Graph(const Graph& copy)
:	m_HeadNode(NULL),
	m_TailNode(NULL)
{
	*this = copy;
}

Graph& Graph::operator=(const Graph& rhs) {
	
	if (this == &rhs) {
		return *this;
	}
	
	/* dealloc any previously allocated dynamic memory */
	freeMemory();
	m_Nodes.clear();
	
	try {
		/* copy values */
		m_HeadNode = rhs.m_HeadNode;
		m_TailNode = rhs.m_TailNode;
		m_Nodes = rhs.m_Nodes;
		
		/* copy edges.
		 * Use the input edge map because this contains only the active edges.
		 * There may be unused edges */
		set<Edge*>::const_iterator edgeIter;
		for (edgeIter = rhs.m_Edges.begin(); edgeIter != rhs.m_Edges.end();
			 edgeIter++) {
			connect(**edgeIter);
		}
	}
	catch (...) {
		/* reset to initial state */
		m_HeadNode = NULL;
		m_TailNode = NULL;
		m_Nodes.clear();
		freeMemory();
		throw;
	}

	return *this;
}

Graph::~Graph() throw() {
	try {
		freeMemory();
	}
	catch (...) {
	}
}

void Graph::setHeadNode(Node& node) {
	m_HeadNode = &node;
	m_Nodes.insert(&node);
}

void Graph::setTailNode(Node& node) {
	m_TailNode = &node;
	m_Nodes.insert(&node);
}

void Graph::addNode(Node& node) {
	m_Nodes.insert(&node);
}

void Graph::connect(const Edge& edge) {
	auto_ptr<Edge> newEdge(new Edge(edge));
	Edge* pEdge = newEdge.get();
	
	/* TODO: Ensure that an edge between the two nodes doesn't already exist,
	 * and if it does, replace it */
	try {
		if (edge.to != NULL) {
			m_Nodes.insert(edge.to);
			
			if (m_InputEdges.count(edge.to) == 0) {
				m_InputEdges[edge.to] = new vector<Edge*>();
			}
			
			m_InputEdges[edge.to]->push_back(newEdge.get());
		}
		
		if (edge.from != NULL) {
			m_Nodes.insert(edge.from);
			
			if (m_OutputEdges.count(edge.from) == 0) {
				m_OutputEdges[edge.from] = new vector<Edge*>();
			}
			
			m_OutputEdges[edge.from]->push_back(newEdge.get());
		}
		
		m_Edges.insert(newEdge.release());
	}
	catch (...) {
		/* need to cleanup so that we don't reference a null ptr later on */
		/* revove from input edges */
		if (m_InputEdges.count(edge.to) != 0) {
			vector<Edge*>* edgeVec = m_InputEdges[edge.to];
			vector<Edge*>::iterator edgeIter;
			edgeIter = find(edgeVec->begin(), edgeVec->end(), pEdge);
			if (edgeIter != edgeVec->end()) {
				edgeVec->erase(edgeIter);
			}
		}
		
		/* remove from output edges */
		if (m_OutputEdges.count(edge.from) != 0) {
			vector<Edge*>* edgeVec = m_OutputEdges[edge.from];
			vector<Edge*>::iterator edgeIter;
			edgeIter = find(edgeVec->begin(), edgeVec->end(), pEdge);
			if (edgeIter != edgeVec->end()) {
				edgeVec->erase(edgeIter);
			}
		}
		
		/* remove from edge set */
		m_Edges.erase(pEdge);
		
		/* make sure it gets deleted */
		if (newEdge.get() == NULL) {
			delete pEdge;
			pEdge = NULL;
		}
		
		throw;
	}
}

void Graph::removeNode(Node& node) {
	m_Nodes.erase(&node);
	
	/* erase node from edge maps and deallocate vectors of edges. */
	if (m_OutputEdges.count(&node) != 0) {
		freeEdges(*m_OutputEdges[&node]);
		delete m_OutputEdges[&node];
		m_OutputEdges.erase(&node);
	}
	
	if (m_InputEdges.count(&node) != 0) {
		freeEdges(*m_InputEdges[&node]);
		delete m_InputEdges[&node];
		m_InputEdges.erase(&node);
	}
}

bool Graph::disconnect(Node* from, Node* to) {
	Edge tempEdge(from, FrameLink(), to);
	Edge* pEdge = NULL;
	
	/* remove edge from input edge mapping */
	if (m_InputEdges.count(to) != 0) {
		vector<Edge*>& edgeVec = *m_InputEdges[to];
		vector<Edge*>::iterator edgeIter;
		edgeIter = search_n(edgeVec.begin(), edgeVec.end(), 
							1, &tempEdge, 
							&equalNodes);
		
		if (edgeIter != edgeVec.end()) {
			edgeVec.erase(edgeIter);
			freeEdge(*edgeIter);
			pEdge = *edgeIter;
		}
	}
	
	/* remove edge from output edge mapping */
	if (m_OutputEdges.count(from) != 0) {
		vector<Edge*>& edgeVec = *m_OutputEdges[from];
		vector<Edge*>::iterator edgeIter;
		edgeIter = search_n(edgeVec.begin(), edgeVec.end(), 
							1, &tempEdge, 
							&equalNodes);
		
		if (edgeIter != edgeVec.end()) {
			edgeVec.erase(edgeIter);
			freeEdge(*edgeIter);
			pEdge = *edgeIter;
		}
	}
	
	return (pEdge != NULL);
}	

Node* Graph::getHeadNode() {
	return m_HeadNode;
}

const Node* Graph::getHeadNode() const {
	return m_HeadNode;
}

Node* Graph::getTailNode() {
	return m_TailNode;
}

const Node* Graph::getTailNode() const {
	return m_TailNode;
}

const set<Node*>& Graph::getNodes() const {
	return m_Nodes;
}

vector<Edge*> Graph::getOutputEdges(const Node* node) const {
	if (m_OutputEdges.count(node) != 0) {
		return *(m_OutputEdges.find(node)->second);
	}
	
	return vector<Edge*>();
}

vector<Edge*> Graph::getInputEdges(const Node* node) const {
	if (m_InputEdges.count(node) != 0) {
		return *m_InputEdges.find(node)->second;
	}
	
	return vector<Edge*>();	
}

void Graph::freeMemory() {

	try {
		/* dealloc edges */
		set<Edge*>::iterator edgeIter;
		for (edgeIter = m_Edges.begin(); edgeIter != m_Edges.end(); 
			 edgeIter++) {
			delete *edgeIter;
		}
		m_Edges.clear();

		/* delete mapping vectors */
		map<const Node*, vector<Edge*>*>::iterator mapIter;
		for (mapIter = m_InputEdges.begin(); mapIter != m_InputEdges.end();
			 mapIter++) {
			delete mapIter->second;
		}
		m_InputEdges.clear();

		for (mapIter = m_OutputEdges.begin(); mapIter != m_OutputEdges.end();
			 mapIter++) {
			delete mapIter->second;
		}
		m_OutputEdges.clear();
	}
	catch (...) {
		/* if exception thrown, allow memory leak but make sure we won't 
		 * attempt to address something that's been deallocated */
		m_Edges.clear();
		m_InputEdges.clear();
		m_OutputEdges.clear();
		throw;
	}
}

void Graph::freeEdges(vector<Edge*>& edgeVec) {
	vector<Edge*>::iterator edgeIter;
	for (edgeIter = edgeVec.begin(); edgeIter != edgeVec.end(); edgeIter++) {
		freeEdge(*edgeIter);
	}
	edgeVec.clear();
}

bool Graph::freeEdge(Edge* edge) {
	if (m_Edges.erase(edge) != 0) {
		delete edge;
		return true;
	}
	
	return false;
}

bool equalNodes(const Edge* edge1, const Edge* edge2) {
	if (edge1 != NULL && edge2 != NULL) {
		if (edge1->to != NULL && edge1->from != NULL && 
			edge2->to != NULL && edge2->from != NULL) 
		{
			return (edge1->to == edge2->to) && (edge1->from == edge2->from);
		}
	}
	return false;
}
